#include "encode_h264_aac.h"
#include <iostream>
ENCODE_H264_AAC::ENCODE_H264_AAC() 
    :codec_ctx_(nullptr),
    av_frame_(nullptr),
    sws_ctx_(nullptr),
    pkt_(nullptr),
    pts_(0),
    width_(0),
    height_(0),
    fps_(0),
    frame_size_(0) {
        
}

ENCODE_H264_AAC::~ENCODE_H264_AAC(){
    cleanup();
}

bool ENCODE_H264_AAC::init_h264(int width, int height, int fps, int bitrate) {
    width_ = width;
    height_ = height;
    fps_ = fps;

    //查找编码器
    const AVCodec *codec = avcodec_find_encoder_by_name("h264_rkmpp");
    if (!codec) {
        std::cerr << "找不到h264_rkmpp编码器" << std::endl;
        return false;
    } else {
        std::cerr << "使用h264_rkmpp编码器" << std::endl;
    }

    // 分配编码上下文
    codec_ctx_ = avcodec_alloc_context3(codec);
    if (!codec_ctx_) {
        std::cerr << "无法分配编码上下文" << std::endl;
        return false;
    }

    //设置编码器参数
    codec_ctx_->bit_rate = bitrate; //比特率 越大质量越好，但文件更大
    codec_ctx_->width = width_;
    codec_ctx_->height = height_;
    //定义时间单位，表示每个时间刻度代表多少秒
    // 示例：{1, 30} 表示每个时间单位是 1/30 秒
    codec_ctx_->time_base = (AVRational){1,fps_}; //时间基 fps=30 ->(1,30)
    codec_ctx_->framerate = (AVRational){fps_,1}; //帧率 fps=30 ->(30,1)
    codec_ctx_->gop_size = fps_; //关键帧间隔 即两个I帧的间隔
    codec_ctx_->max_b_frames = 0;//不使用B帧
    codec_ctx_->delay = 0; //禁用编码器缓冲 禁用缓冲，实现最低延迟
    codec_ctx_->flags |= AV_CODEC_FLAG_LOW_DELAY; //启用低延迟模式
    //设置编码器输入的像素格式
    codec_ctx_->pix_fmt = AV_PIX_FMT_YUV420P; //使用YUV420

    /*******************************************提高质量***********************************************/
    // codec_ctx_->bit_rate = 4000000;  // 增加到4Mbps
    // codec_ctx_->gop_size = fps_ * 2; // 关键帧间隔加倍
    /*******************************************提高质量***********************************************/

    /*******************************************降低延迟***********************************************/
    // codec_ctx_->gop_size = fps_ / 2; // 关键帧间隔加倍
    /*******************************************降低延迟***********************************************/

    //设置编码器私有参数 
    av_opt_set(codec_ctx_->priv_data, "preset", "fast", 0); // 编码速度预设 平衡编码速度和质量
    av_opt_set(codec_ctx_->priv_data, "profile", "high", 0); // 定义编码功能的子集，影响兼容性和效率
    //设置GOP（Group of Pictures）大小
    av_opt_set_int(codec_ctx_->priv_data, "gop", fps_, 0); // 每 fps_ 帧插入一个关键帧（I帧）
    //控制编码质量的核心参数
    av_opt_set_int(codec_ctx_->priv_data, "qp", 26, 0); // 与比特率的关系：QP值越小，质量越好，但比特率越高

    //打开编码器
    if (avcodec_open2(codec_ctx_, codec, NULL) < 0) {
        std::cerr << "无法打开编码器" << std::endl;
        avcodec_free_context(&codec_ctx_);
        return false;
    }

    //分配 AVFrame
    av_frame_ = av_frame_alloc();
    if (!av_frame_) {
        std::cerr << "无法分配AVFrame" << std::endl;
        cleanup();
        return false;
    }

    av_frame_->format = codec_ctx_->pix_fmt;
    av_frame_->width   = codec_ctx_->width;
    av_frame_->height = codec_ctx_->height;
    //分配数据缓冲区 使用32字节对齐
    if (av_frame_get_buffer(av_frame_ , 32) < 0) {
        std::cerr << "无法分配帧缓冲区" << std::endl;
        cleanup();
        return false;
    }

    // 设置图像缩放转换上下文
    // sws_ctx_ = sws_getContext(640, 480, AV_PIX_FMT_UYVY422,
    //                          width_, height_, AV_PIX_FMT_YUV420P,
    //                          SWS_BILINEAR, NULL, NULL, NULL);
    // if (!sws_ctx_) {
    //     std::cerr << "无法创建缩放上下文" << std::endl;
    //     cleanup();
    //     return false;
    // }
    
    //分配包
    pkt_ = av_packet_alloc();
    if (!pkt_) {
        std::cerr << "无法分配AVPacket" << std::endl;
        cleanup();
        return false;
    }

    return true;
}

bool ENCODE_H264_AAC::init_aac(int sample_rate, int channels, int bitrate) {
    sample_rate_ = sample_rate;
    channels_    = channels;
    bit_rate_    = bitrate;

    //查找编码器
    const AVCodec *codec = avcodec_find_encoder(AV_CODEC_ID_AAC);
    if (!codec) {
        std::cerr << "找不到AAC编码器" << std::endl;
        return false;
    } else {
        std::cerr << "使用AAC编码器: " << codec->name << std::endl;
    }

    //分配编码上下文
    codec_ctx_ = avcodec_alloc_context3(codec);
    if (!codec_ctx_) {
        std::cerr << "无法分配编码上下文" << std::endl;
        return false;
    }
    std::cerr << "1" << std::endl;

    //设置编码器参数
    codec_ctx_->bit_rate       = bit_rate_;
    codec_ctx_->sample_rate    = sample_rate_;
    codec_ctx_->channels       = channels_;
    //设置音频的声道布局。                 单声道                立体声
    codec_ctx_->channel_layout = (channels_ == 1) ? AV_CH_LAYOUT_MONO : AV_CH_LAYOUT_STEREO;
    //设置音频的样本格式。          表示浮点平面格式
    codec_ctx_->sample_fmt      = AV_SAMPLE_FMT_FLTP;
    //设置时间基
    codec_ctx_->time_base      =(AVRational){1,sample_rate_};
    // 设置对标准的严格遵守程度。
    // codec_ctx_->strict_std_compliance = FF_COMPLIANCE_EXPERIMENTAL;
    // 如果目标是文件存储
    codec_ctx_->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;
    // 如果目标是流媒体传输
    // codec_ctx_->flags &= ~AV_CODEC_FLAG_GLOBAL_HEADER;

    std::cerr << "2" << std::endl;

    //打开编码器
    if (avcodec_open2(codec_ctx_, codec, nullptr) < 0) {
        std::cerr << "无法打开AAC编码器" << std::endl;
        avcodec_free_context(&codec_ctx_);
        return false;
    }

    frame_size_ = codec_ctx_->frame_size;
    std::cout << "size: " << frame_size_ <<std::endl;

    if (frame_size_ == 0) {
        frame_size_ = 1024;
    }

    std::cout << "AAC编码器初始化成功, 帧大小: " << frame_size_ 
              << ", 采样率: " << sample_rate_ 
              << ", 声道数: " << channels_ << std::endl;

    //分配AVFrame
    av_frame_ = av_frame_alloc();
    if (!av_frame_) {
        std::cerr << "无法分配AVFrame" << std::endl;
        cleanup();
        return false;
    }

    //设置音频帧中的样本数量
    av_frame_->nb_samples     = frame_size_;
    //设置音频帧中的样本格式
    av_frame_->format         = codec_ctx_->sample_fmt;
    //设置音频帧的声道布局。
    av_frame_->channel_layout = codec_ctx_->channel_layout;
    //设置音频帧的采样率
    av_frame_->sample_rate    = sample_rate_;

    //分配帧缓冲区
    if (av_frame_get_buffer(av_frame_, 0) < 0) {
        std::cerr << "无法分配帧缓冲区" << std::endl;
        cleanup();
        return false;
    }
    
    //分配包
    pkt_ = av_packet_alloc();
    if (!pkt_) {
        std::cerr << "无法分配AVPacket" << std::endl;
        cleanup();
        return false;
    }

    // 预分配浮点转换缓冲区
    float_buffer_.resize(frame_size_ * channels_);

    return true;
}

void ENCODE_H264_AAC::add_adts_header(uint8_t* buffer, int packet_size, int profile, int sample_rate, int channels) {
    int sampling_frequency_index = 4; // 默认44100Hz
    switch(sample_rate) {
        case 96000: sampling_frequency_index = 0; break;
        case 88200: sampling_frequency_index = 1; break;
        case 64000: sampling_frequency_index = 2; break;
        case 48000: sampling_frequency_index = 3; break;
        case 44100: sampling_frequency_index = 4; break;
        case 32000: sampling_frequency_index = 5; break;
        case 24000: sampling_frequency_index = 6; break;
        case 22050: sampling_frequency_index = 7; break;
        case 16000: sampling_frequency_index = 8; break;
        case 12000: sampling_frequency_index = 9; break;
        case 11025: sampling_frequency_index = 10; break;
        case 8000: sampling_frequency_index = 11; break;
    }
    
    int frame_length = packet_size + 7;
    
    buffer[0] = 0xFF;
    buffer[1] = 0xF1;
    buffer[2] = ((profile - 1) << 6) | (sampling_frequency_index << 2) | (channels >> 2);
    buffer[3] = ((channels & 3) << 6) | (frame_length >> 11);
    buffer[4] = (frame_length >> 3) & 0xFF;
    buffer[5] = ((frame_length & 7) << 5) | 0x1F;
    buffer[6] = 0xFC;
}

bool ENCODE_H264_AAC::encodeH264Frame(uint8_t * yuv420_data,  std::function<void(AVPacket *)> packet_callback ){
    int ret;
    if (!codec_ctx_ || !av_frame_ ) {
        std::cerr << "编码器未初始化" << std::endl;
        return false;
    } 
    if (!yuv420_data) {
        std::cerr << "输入数据为空" << std::endl;
        return false;
    }

    // 计算YUV420数据大小
    int y_size = width_ * height_;
    //4分y分量等于1分uv分量
    int uv_size = y_size / 4; 
    
    // 复制Y分量
    memcpy(av_frame_->data[0], yuv420_data, y_size);
    // 复制U分量
    memcpy(av_frame_->data[1], yuv420_data + y_size, uv_size);
    // 复制V分量
    memcpy(av_frame_->data[2], yuv420_data + y_size + uv_size, uv_size);

    av_frame_->pts = pts_++;

    /******************************************发帧 取包******************************************/
    // 发生帧到编码器
    if (avcodec_send_frame(codec_ctx_, av_frame_) < 0) {
        std::cerr << "发送帧到编码器失败: " << std::endl;
        return false;
    }

    // 接受编码后的包
    while (1) {
        ret = avcodec_receive_packet(codec_ctx_, pkt_);
        if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) {
            break;
        } else if (ret < 0) {
            std::cerr << "编码错误: " << ret << std::endl;
            break;
        }

        if (packet_callback) {
            packet_callback(pkt_);
        }
        // 因为 传入了lambda表达式 [inter_timestamp_ms = video_frame.timestamp_ms](AVPacket* pkt)
        // 所以 这里的else是不会运行的
        else {
            // 默认处理：克隆包并放入队列
            AVPacket* cloned_pkt = av_packet_clone(pkt_);
            std::lock_guard<std::mutex> lock(queue_mutex_);
            packet_queue_.push(cloned_pkt);
            queue_cv_.notify_one();
        }
        
        av_packet_unref(pkt_);
    }    
    /******************************************发帧 取包******************************************/

    return true;
}

bool ENCODE_H264_AAC::encodeAACFrame(uint16_t * pcm_data, int samples ,std::function<void(AVPacket *)> packet_callback ){
    if (!codec_ctx_ || !av_frame_) {
        std::cerr << "编码器未初始化" << std::endl;
        return false;
    } 
    if (!pcm_data || samples <= 0) {
        std::cerr << "输入数据无效" << std::endl;
        return false;
    }


    // 准备音频帧
    if (av_frame_make_writable(av_frame_) < 0) {
        std::cerr << "无法使帧可写" << std::endl;
        return false;
    }

    // 将16位整型PCM转换为浮点平面格式
    for (int i = 0; i < frame_size_; i++) {
        for (int ch = 0; ch < channels_; ch++) {
            float *dst = (float *)av_frame_->data[ch];
            // 从交错格式中提取每个声道的样本
            int16_t sample = pcm_data[i * channels_ + ch];
            // 将16位有符号整数转换为[-1.0, 1.0]范围的浮点数
            dst[i] = sample / 32768.0f;
            
            // 添加限制，防止溢出
            if (dst[i] > 1.0f) dst[i] = 1.0f;
            if (dst[i] < -1.0f) dst[i] = -1.0f;
        }
    }
    
    av_frame_->pts = pts_;
    pts_ += frame_size_;

    int ret = avcodec_send_frame(codec_ctx_, av_frame_);
    if (ret < 0) {
        std::cerr << "发送帧到编码器失败: " << ret << std::endl;
        return false;
    }

    while(1) {
        ret = avcodec_receive_packet (codec_ctx_, pkt_);
        if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) {
            break;
        } else if (ret < 0) {
            std::cerr << "编码错误: " << ret << std::endl;
            return false;
        }

        if (packet_callback) {
            packet_callback(pkt_);
        } 
        // 因为 传入了lambda表达式 [inter_timestamp_ms = video_frame.timestamp_ms](AVPacket* pkt)
        // 所以 这里的else是不会运行的
        else {
            // 默认处理：克隆包并放入队列
            AVPacket* cloned_pkt = av_packet_clone(pkt_);
            std::lock_guard<std::mutex> lock(queue_mutex_);
            packet_queue_.push(cloned_pkt);
            queue_cv_.notify_one();
        }
        av_packet_unref(pkt_);
    }
    return true;
}


AVPacket* ENCODE_H264_AAC::getH264Packet(){
    std::unique_lock<std::mutex> lock(queue_mutex_);
    if (packet_queue_.empty()) {
        // 等待一段时间看看是否有包到来
        if (queue_cv_.wait_for(lock, std::chrono::milliseconds(100)) == std::cv_status::timeout) {
            std::cout << "包队列为空" << std::endl;
            return nullptr;
        }
    }

    AVPacket *pkt = packet_queue_.front();
    packet_queue_.pop();
    return pkt;
}

AVPacket* ENCODE_H264_AAC::getAACPacket() {
    std::unique_lock<std::mutex> lock(queue_mutex_);
    if (packet_queue_.empty()) {
        // 等待一段时间看看是否有包到来
        if (queue_cv_.wait_for(lock, std::chrono::milliseconds(100)) == std::cv_status::timeout) {
            return nullptr;
        }
    }

    if (packet_queue_.empty()) {
        return nullptr;
    }

    AVPacket *pkt = packet_queue_.front();
    packet_queue_.pop();
    return pkt;
}

AVRational ENCODE_H264_AAC::get_time_base() const { 
    return codec_ctx_ ? codec_ctx_->time_base : AVRational{0, 1}; 
}

bool ENCODE_H264_AAC::flushEncoder(std::function<void(AVPacket *)> packet_callback) {
    int ret;
    
    // 发送null帧来刷新编码器
    ret = avcodec_send_frame(codec_ctx_, nullptr);
    if (ret < 0) {
        std::cerr << "刷新编码器失败: " << ret << std::endl;
        return false;
    }

    // 接收所有剩余的包
    while (true) {
        ret = avcodec_receive_packet(codec_ctx_, pkt_);
        if (ret == AVERROR_EOF) {
            break;
        } else if (ret < 0) {
            std::cerr << "接收包错误: " << ret << std::endl;
            return false;
        }

        if (packet_callback) {
            packet_callback(pkt_);
        } else {
            AVPacket* cloned_pkt = av_packet_clone(pkt_);
            std::lock_guard<std::mutex> lock(queue_mutex_);
            packet_queue_.push(cloned_pkt);
            queue_cv_.notify_one();
        }
        
        av_packet_unref(pkt_);
    }

    return true;
}

void ENCODE_H264_AAC::cleanup() {
    if (pkt_) {
        av_packet_free(&pkt_);
        pkt_ = nullptr;
    }
    
    if (av_frame_) {
        av_frame_free(&av_frame_);
        av_frame_ = nullptr;
    }
    
    if (codec_ctx_) {
        avcodec_free_context(&codec_ctx_);
        codec_ctx_ = nullptr;
    }
    
    // if (sws_ctx_) {
    //     sws_freeContext(sws_ctx_);
    //     sws_ctx_ = nullptr;
    // }
    
    // 清理包队列
    std::lock_guard<std::mutex> lock(queue_mutex_);
    while (!packet_queue_.empty()) {
        AVPacket* pkt = packet_queue_.front();
        av_packet_free(&pkt);
        packet_queue_.pop();
    }
    pts_ = 0;
    sample_rate_ = 0;
    channels_ = 0;
    bit_rate_ = 0;
    frame_size_ = 0;
    float_buffer_.clear();
}

int ENCODE_H264_AAC::getFrameSize() {
    return frame_size_;
}